import os, json
import numpy as np
import pandas as pd
from pyproj import Transformer
from scipy.spatial.transform import Slerp
from scipy.spatial.transform import Rotation
from open3d import geometry

'''

'''
UTC_ZONE = 8


def umeyama(X, Y):
    """
    Estimates the Sim(3) transformation between `X` and `Y` point sets.
    Estimates c, R and t such as c * R @ X + t ~ Y.
    Parameters
    ----------
    X : numpy.array
        (m, n) shaped numpy array. m is the dimension of the points,
        n is the number of points in the point set.
    Y : numpy.array
        (m, n) shaped numpy array. Indexes should be consistent with `X`.
        That is, Y[:, i] must be the point corresponding to X[:, i].

    Returns
    -------
    c : float
        Scale factor.
    R : numpy.array
        (3, 3) shaped rotation matrix.
    t : numpy.array
        (3, 1) shaped translation vector.
    """
    mu_x = X.mean(axis=1).reshape(-1, 1)
    mu_y = Y.mean(axis=1).reshape(-1, 1)
    var_x = np.square(X - mu_x).sum(axis=0).mean()
    cov_xy = ((Y - mu_y) @ (X - mu_x).T) / X.shape[1]
    U, D, VH = np.linalg.svd(cov_xy)
    S = np.eye(X.shape[0])
    if np.linalg.det(U) * np.linalg.det(VH) < 0:
        S[-1, -1] = -1
    c = np.trace(np.diag(D) @ S) / var_x
    R = U @ S @ VH
    t = mu_y - c * R @ mu_x
    return c, R, t


def rt2homo(R, t):
    return np.block([[R, t.reshape(3,1)], [np.zeros((1, 3)), 1]])


# calculate the rotation matrix from Euler angel defined by our system
def get_euler_matrix(yaw, pitch, roll):
    return geometry.get_rotation_matrix_from_zxy(np.array([-yaw, pitch, roll]) * np.pi / 180)


# rotation matrix from ECEF to local ENU
def get_rot_ecef_enu(lat, lon):
    return geometry.get_rotation_matrix_from_zyx(np.array([90 + lon, 0, 90 - lat]) * np.pi / 180)


class POSFile:
    def __init__(self, pos_path, conf_path):
        # read POS file
        # regard POS file as a space seperated table
        # (except first 23 lines and line 25, which are some comments)
        self.pos_df = pd.read_table(pos_path, sep='\s+', skiprows=list(range(23)) + [24])
        # self.conf = json.load(conf_path)
        # WGS-84 to projection
        transformer_proj = Transformer.from_crs("epsg:4326", "epsg:32648")
        # Geodetic to Geocentric
        transformer_ecef = Transformer.from_crs(
            "EPSG:4326",
            {"proj": 'geocent', "ellps": 'WGS84', "datum": 'WGS84'}
        )

        self.proj_xyz = np.array(transformer_proj.transform(self.pos_df['Latitude'],
                                                            self.pos_df['Longitude']
                                                            ) + tuple([self.pos_df["H-Ell"]]))
        self.ecef_xyz = np.array(transformer_ecef.transform(self.pos_df['Latitude'],
                                                            self.pos_df['Longitude'],
                                                            self.pos_df["H-Ell"]))
        # camera setting
        self.T_b_c = np.array([-0.058, 0.005, 0.8891])
        self.C_b_c = get_euler_matrix(89.4, -0.1, 0.2)

    def get_rot_base_ecef(self, idx):
        pos_t = self.pos_df.iloc[idx]
        C_w_l = get_rot_ecef_enu(pos_t["Latitude"], pos_t['Longitude'])
        C_l_b = get_euler_matrix(pos_t["Heading"], pos_t["Pitch"], pos_t["Roll"])
        return C_w_l @ C_l_b

    def get_time_idx(self, time):
        i = 0
        while i < len(self.pos_df) and self.pos_df.loc[i]["UTCTime"] <= time:
            i += 1
        dt = time - self.pos_df.loc[i - 1]["UTCTime"]
        return i - 1, dt / (self.pos_df.loc[i]["UTCTime"] - self.pos_df.loc[i - 1]["UTCTime"])

    def get_model_ecef(self):
        c, R, t = umeyama(self.proj_xyz, self.ecef_xyz)
        return c * R, t

    def get_base_ecef_by_time(self, time):
        i0, p = self.get_time_idx(time)
        slerp = Slerp([0, 1], Rotation.from_matrix([self.get_rot_base_ecef(i0), self.get_rot_base_ecef(i0 + 1)]))
        C_w_b = slerp(p).as_matrix()
        # pos_t = self.pos_df.iloc[i0]
        # C_w_l = get_rot_ecef_enu(pos_t['Longitude'], pos_t["Latitude"])
        # v = np.array(pos_t[["VEast", "VNorth", "VUp"]])
        # a = np.array(pos_t[["AccEast", "AccNrth", "AccUp"]])
        T_w_b = self.ecef_xyz[:, i0] * (1 - p) + self.ecef_xyz[:, i0 + 1] * p
        return C_w_b, T_w_b

    def get_cam_ecef_by_time(self, time):
        C_w_b, T_w_b = self.get_base_ecef_by_time(time)
        T_w_c = T_w_b + C_w_b @ self.T_b_c  # camera in world
        C_w_c = C_w_b @ self.C_b_c
        return C_w_c, T_w_c


if __name__ == "__main__":
    POS_path = "/home/ivan/Datasets/SRTP-DATA/20210112-CoreSixPano/CoreSix.pos"
    a = POSFile(POS_path, "")
    print(a.get_base_ecef_by_time(9811.032))
    print(a.get_model_ecef())
